package com.example.mybatisplusdemo.controller;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.alibaba.excel.write.style.column.LongestMatchColumnWidthStyleStrategy;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.example.mybatisplusdemo.entity.User;
import com.example.mybatisplusdemo.mapper.UserMapper;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.List;

/**
 * @author 江黎
 * @since 2023-02-24
 */
@RestController
@RequestMapping
public class ExportController {
    @Resource
    UserMapper userMapper;

    public static final Integer PER_SHEET_ROW_COUNT = 100000;
    public static final Integer PER_WRITE_ROW_COUNT = 100000;

    @PostMapping("/export")
    public void export(HttpServletResponse response) {
        try (OutputStream outputStream = response.getOutputStream();) {
            long startTime = System.currentTimeMillis();
            System.out.println("导出开始时间:" + startTime);

            //记录总数:实际中需要根据查询条件进行统计即可:一共多少条
            Integer totalCount = Math.toIntExact(userMapper.selectCount(null));


            //每一个Sheet存放的条数: 10000
            Integer sheetDataRows = PER_SHEET_ROW_COUNT;
            //每次写入的数据量: 10000
            Integer writeDataRows = PER_WRITE_ROW_COUNT;
            //计算需要的Sheet数量
            Integer sheetNum = totalCount % sheetDataRows == 0 ? (totalCount / sheetDataRows) : (totalCount / sheetDataRows + 1);
            //计算一般情况下每一个Sheet需要写入的次数(一般情况不包含最后一个sheet,因为最后一个sheet不确定会写入多少条数据)
            Integer oneSheetWriteCount = sheetDataRows / writeDataRows;
            //计算最后一个sheet需要写入的次数
            Integer lastSheetWriteCount = totalCount % sheetDataRows == 0 ? oneSheetWriteCount : (totalCount % sheetDataRows % writeDataRows == 0 ? (totalCount / sheetDataRows / writeDataRows) : (totalCount / sheetDataRows / writeDataRows + 1));

            //必须放到循环外，否则会刷新流
            ExcelWriter excelWriter = EasyExcel.write(outputStream).build();

            List<User> dataList = new ArrayList<>();
            //开始分批查询分次写入
            for (int i = 0; i < sheetNum; i++) {
                //创建Sheet
                WriteSheet sheet = new WriteSheet();
                sheet.setSheetName("Sheet" + i);
                //循环写入次数: j的自增条件是当不是最后一个Sheet的时候写入次数为正常的每个Sheet写入的次数,如果是最后一个就需要使用计算的次数lastSheetWriteCount
                for (int j = 0; j < (i != sheetNum - 1 ? oneSheetWriteCount : lastSheetWriteCount); j++) {
                    //集合复用,便于GC清理
                    dataList.clear();
                    //分页查询一次1000
                    Page<User> page = new Page<>(j + 1 + oneSheetWriteCount * i, writeDataRows);
                    page = userMapper.selectPage(page, null);
                    dataList = page.getRecords();
                    //写数据
                    WriteSheet writeSheet = EasyExcel.writerSheet(i, "Sheet" + (i + 1))
                            .registerWriteHandler(new LongestMatchColumnWidthStyleStrategy()).build();
                    excelWriter.write(dataList, writeSheet);
                }
            }

            // 下载EXCEL
            response.setContentType("application/octet-stream");
            response.setCharacterEncoding("utf-8");
            response.setHeader("Content-Disposition", "attachment;filename=" + new String(("导出结果").getBytes("gb2312"), StandardCharsets.ISO_8859_1) + ".xlsx");
            excelWriter.finish();
            outputStream.flush();

            //导出时间结束
            long endTime = System.currentTimeMillis();
            System.out.println("导出结束时间:" + endTime + "ms");
            System.out.println("导出所用时间:" + (endTime - startTime) / 1000 + "秒");
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
